from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware

from api.routers.chat import router as chat_router
from api.routers.documents import router as documents_router
from api.routers.knowledge_base import router as kb_router
from api.routers.search import router as search_router
from api.routers.system_prompt import router as sp_router
from api.routers.statistics import router as statistics_router
from api.routers.models import router as models_router
from api.routers.target_database import router as target_db_router
from api.routers.auth import router as auth_router


app = FastAPI(
    title="AI cang",
    description="AI助手系统API",
    version="1.0.0",
    # 添加Swagger UI的认证配置
    swagger_ui_parameters={
        "persistAuthorization": True,
    }
)

# 配置安全方案
app.openapi_schema = None

def custom_openapi():
    if app.openapi_schema:
        return app.openapi_schema
    
    from fastapi.openapi.utils import get_openapi
    
    openapi_schema = get_openapi(
        title=app.title,
        version=app.version,
        description=app.description,
        routes=app.routes,
    )
    
    # 添加Bearer认证方案
    openapi_schema["components"]["securitySchemes"] = {
        "BearerAuth": {
            "type": "http",
            "scheme": "bearer",
            "bearerFormat": "JWT",
            "description": "请输入JWT token，格式：Bearer <token>"
        }
    }
    
    # 为需要认证的路径添加安全要求
    for path, path_info in openapi_schema.get("paths", {}).items():
        for method, method_info in path_info.items():
            # 为除了登录注册外的认证相关接口添加安全要求
            if path.startswith("/auth/") and path not in ["/auth/login", "/auth/register"]:
                method_info["security"] = [{"BearerAuth": []}]
    
    app.openapi_schema = openapi_schema
    return app.openapi_schema

app.openapi = custom_openapi

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,  # 修改为True以支持认证
    allow_methods=["*"],
    allow_headers=["*"],
)

# 路由分组挂载（保持原有路径不变）
app.include_router(auth_router, prefix="/auth", tags=["认证"])
app.include_router(chat_router)
app.include_router(documents_router)
app.include_router(kb_router)
app.include_router(search_router)
app.include_router(sp_router)
app.include_router(statistics_router)
app.include_router(models_router)
app.include_router(target_db_router)


